# Chat/model_handler.py
from typing import Dict, Any

from langchain_core.runnables import RunnablePassthrough, RunnableParallel
from langchain_core.output_parsers import StrOutputParser


class Text2SQLGenerator:
    """
    支持上下文的Text2SQL生成器
    Attributes:
          llm: 用于生成的LLM模型
          prompt_template: 用于生成的PromptTemplate
          memory: 用于生成的Memory
    """

    def __init__(self, llm, prompt_template, memory):
        self.llm = llm
        self.prompt_template = prompt_template
        self.memory = memory

        # 构建新的链式结构
        self.chain = (
                RunnableParallel(
                    # 动态获取当前输入参数和内存中的历史记录
                    dialect=lambda x: x["dialect"],
                    table_info=lambda x: x["table_info"],
                    question=lambda x: x["question"],
                    history=lambda _: self.memory.load_memory_variables({})["history"]
                )
                | self.prompt_template
                | self.llm
                | StrOutputParser()
        )

    def generate_sql(self, inputs: Dict[str, Any]) -> str:
        """生成并返回SQL语句
        Args:
            inputs: 输入参数，包括dialect, table_info, question
        Returns:
            str: 生成的SQL语句
        """
        response = self.chain.invoke(inputs)
        return response.strip()


class SQL2TextGenerator:
    """
    支持上下文的SQL2Text生成器
    Attributes:
        llm: 用于生成的LLM模型
        prompt_template: 用于生成的PromptTemplate
    """

    def __init__(self, llm, prompt_template):
        self.llm = llm
        self.prompt_template = prompt_template

        self.chain = (
                RunnableParallel(
                    result=lambda x: x["result"],
                    question=lambda x: x["question"],
                    sqlcode=lambda x: x["sqlcode"],
                )
                | self.prompt_template
                | self.llm
                | StrOutputParser()
        )

    def generate_text(self, inputs: Dict[str, Any]) -> str:
        """
        生成并返回自然语言回答
        :param inputs: 输入参数，包括result, question, sqlcode
        :return: 生成的自然语言回答
        """
        """生成并返回自然语言回答"""
        response = self.chain.invoke(inputs)
        return response.strip()
